import numba as nb
import numpy as np
import torch

from utils.compressor import Compressor


class Sparser(Compressor):
    def __init__(self, ratio: float = 0.98, *args, **kwargs) -> None:
        self.ratio = ratio
        self.arg = args
        self.kwargs = kwargs

    def __str__(self):
        return f"Top-{int((1 - self.ratio) * 100)}% sparsification"

    def encode(self, x: torch.Tensor) -> (torch.Tensor, torch.Tensor):
        with torch.no_grad():
            order = abs(x.data.view(-1)).sort()[0]
            threshold = float(order[int(self.ratio * len(order)) - 1])
            mask = torch.where(abs(x.data) > threshold, 1.0, 0.0)
            vector = torch.multiply(mask, x.data)
            vector = vector[vector > 0].view(-1)
        return vector, mask

    def decode(self, x: torch.Tensor, vector: torch.Tensor, mask: torch.Tensor) -> None:
        with torch.no_grad():
            mask = mask.view(-1).numpy()
            self.__decode(vector.numpy(), mask)
            x.data = torch.from_numpy(mask).view(x.size())

    @staticmethod
    @nb.jit
    def __decode(vector: np.ndarray, mask: np.ndarray) -> None:
        cur = 0
        for idx, val in enumerate(mask):
            if val:
                mask[idx] = vector[cur]
                cur += 1
